/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.state;



import com.bff.gaia.unified.runners.core.StateInternals;

import com.bff.gaia.unified.runners.core.StateNamespace;

import com.bff.gaia.unified.runners.core.StateTag;

import com.bff.gaia.unified.runners.gaia.translation.types.CoderTypeSerializer;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.CoderException;

import com.bff.gaia.unified.sdk.coders.InstantCoder;

import com.bff.gaia.unified.sdk.coders.VoidCoder;

import com.bff.gaia.unified.sdk.state.BagState;

import com.bff.gaia.unified.sdk.state.CombiningState;

import com.bff.gaia.unified.sdk.state.MapState;

import com.bff.gaia.unified.sdk.state.ReadableState;

import com.bff.gaia.unified.sdk.state.ReadableStates;

import com.bff.gaia.unified.sdk.state.SetState;

import com.bff.gaia.unified.sdk.state.State;

import com.bff.gaia.unified.sdk.state.StateBinder;

import com.bff.gaia.unified.sdk.state.StateContext;

import com.bff.gaia.unified.sdk.state.StateSpec;

import com.bff.gaia.unified.sdk.state.ValueState;

import com.bff.gaia.unified.sdk.state.WatermarkHoldState;

import com.bff.gaia.unified.sdk.transforms.Combine;

import com.bff.gaia.unified.sdk.transforms.CombineWithContext;

import com.bff.gaia.unified.sdk.transforms.windowing.BoundedWindow;

import com.bff.gaia.unified.sdk.transforms.windowing.TimestampCombiner;

import com.bff.gaia.unified.sdk.util.CoderUtils;

import com.bff.gaia.unified.sdk.util.CombineContextFactory;

import com.bff.gaia.unified.vendor.guava.com.google.common.base.Preconditions;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableList;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Iterables;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Lists;

import com.bff.gaia.api.common.state.ListState;

import com.bff.gaia.api.common.state.ListStateDescriptor;

import com.bff.gaia.api.common.state.MapStateDescriptor;

import com.bff.gaia.api.common.state.ValueStateDescriptor;

import com.bff.gaia.api.common.typeutils.base.BooleanSerializer;

import com.bff.gaia.api.common.typeutils.base.StringSerializer;

import com.bff.gaia.runtime.state.KeyedStateBackend;

import com.bff.gaia.runtime.state.VoidNamespace;

import com.bff.gaia.runtime.state.VoidNamespaceSerializer;

import org.joda.time.Instant;



import javax.annotation.Nonnull;

import java.nio.ByteBuffer;

import java.util.Collections;

import java.util.HashMap;

import java.util.Iterator;

import java.util.Map;

import java.util.stream.Stream;



/**

 * {@link StateInternals} that uses a Gaia {@link KeyedStateBackend} to manage state.

 *

 * <p>Note: In the Gaia streaming runner the key is always encoded using an {@link Coder} and

 * stored in a {@link ByteBuffer}.

 */

public class GaiaStateInternals<K> implements StateInternals {



  private final KeyedStateBackend<ByteBuffer> gaiaStateBackend;

  private Coder<K> keyCoder;



  // Combined watermark holds for all keys of this partition

  private final Map<String, Instant> watermarkHolds = new HashMap<>();

  // State to persist combined watermark holds for all keys of this partition

  private final MapStateDescriptor<String, Instant> watermarkHoldStateDescriptor =

      new MapStateDescriptor<>(

          "watermark-holds",

          StringSerializer.INSTANCE,

          new CoderTypeSerializer<>(InstantCoder.of()));



  public GaiaStateInternals(KeyedStateBackend<ByteBuffer> gaiaStateBackend, Coder<K> keyCoder)

      throws Exception {

    this.gaiaStateBackend = gaiaStateBackend;

    this.keyCoder = keyCoder;

    restoreWatermarkHoldsView();

  }



  /** Returns the minimum over all watermark holds. */

  public Instant watermarkHold() {

    long min = Long.MAX_VALUE;

    for (Instant hold : watermarkHolds.values()) {

      min = Math.min(min, hold.getMillis());

    }

    return new Instant(min);

  }



  @Override

  public K getKey() {

    ByteBuffer keyBytes = gaiaStateBackend.getCurrentKey();

    byte[] bytes = new byte[keyBytes.remaining()];

    keyBytes.get(bytes);

    keyBytes.position(keyBytes.position() - bytes.length);

    try {

      return CoderUtils.decodeFromByteArray(keyCoder, bytes);

    } catch (CoderException e) {

      throw new RuntimeException("Error decoding key.", e);

    }

  }



  @Override

  public <T extends State> T state(

      StateNamespace namespace, StateTag<T> address, StateContext<?> context) {

    return address

        .getSpec()

        .bind(

            address.getId(),

            new GaiaStateBinder(

                namespace,

                context,

                gaiaStateBackend,

                watermarkHolds,

                watermarkHoldStateDescriptor));

  }



  private static class GaiaStateBinder implements StateBinder {



    private final StateNamespace namespace;

    private final StateContext<?> stateContext;

    private final KeyedStateBackend<ByteBuffer> gaiaStateBackend;

    private final Map<String, Instant> watermarkHolds;

    private final MapStateDescriptor<String, Instant> watermarkHoldStateDescriptor;



    private GaiaStateBinder(

        StateNamespace namespace,

        StateContext<?> stateContext,

        KeyedStateBackend<ByteBuffer> gaiaStateBackend,

        Map<String, Instant> watermarkHolds,

        MapStateDescriptor<String, Instant> watermarkHoldStateDescriptor) {

      this.namespace = namespace;

      this.stateContext = stateContext;

      this.gaiaStateBackend = gaiaStateBackend;

      this.watermarkHolds = watermarkHolds;

      this.watermarkHoldStateDescriptor = watermarkHoldStateDescriptor;

    }



    @Override

    public <T2> ValueState<T2> bindValue(

		String id, StateSpec<ValueState<T2>> spec, Coder<T2> coder) {

      return new GaiaValueState<>(gaiaStateBackend, id, namespace, coder);

    }



    @Override

    public <T2> BagState<T2> bindBag(String id, StateSpec<BagState<T2>> spec, Coder<T2> elemCoder) {

      return new GaiaBagState<>(gaiaStateBackend, id, namespace, elemCoder);

    }



    @Override

    public <T2> SetState<T2> bindSet(String id, StateSpec<SetState<T2>> spec, Coder<T2> elemCoder) {

      return new GaiaSetState<>(gaiaStateBackend, id, namespace, elemCoder);

    }



    @Override

    public <KeyT, ValueT> MapState<KeyT, ValueT> bindMap(

        String id,

        StateSpec<MapState<KeyT, ValueT>> spec,

        Coder<KeyT> mapKeyCoder,

        Coder<ValueT> mapValueCoder) {

      return new GaiaMapState<>(gaiaStateBackend, id, namespace, mapKeyCoder, mapValueCoder);

    }



    @Override

    public <InputT, AccumT, OutputT> CombiningState<InputT, AccumT, OutputT> bindCombining(

        String id,

        StateSpec<CombiningState<InputT, AccumT, OutputT>> spec,

        Coder<AccumT> accumCoder,

        Combine.CombineFn<InputT, AccumT, OutputT> combineFn) {

      return new GaiaCombiningState<>(gaiaStateBackend, id, combineFn, namespace, accumCoder);

    }



    @Override

    public <InputT, AccumT, OutputT>

        CombiningState<InputT, AccumT, OutputT> bindCombiningWithContext(

            String id,

            StateSpec<CombiningState<InputT, AccumT, OutputT>> spec,

            Coder<AccumT> accumCoder,

            CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn) {

      return new GaiaCombiningStateWithContext<>(

          gaiaStateBackend,

          id,

          combineFn,

          namespace,

          accumCoder,

          CombineContextFactory.createFromStateContext(stateContext));

    }



    @Override

    public WatermarkHoldState bindWatermark(

        String id, StateSpec<WatermarkHoldState> spec, TimestampCombiner timestampCombiner) {

      return new GaiaWatermarkHoldState<>(

          gaiaStateBackend,

          watermarkHolds,

          watermarkHoldStateDescriptor,

          id,

          namespace,

          timestampCombiner);

    }

  }



  private static class GaiaValueState<T> implements ValueState<T> {



    private final StateNamespace namespace;

    private final String stateId;

    private final ValueStateDescriptor<T> gaiaStateDescriptor;

    private final KeyedStateBackend<ByteBuffer> gaiaStateBackend;



    GaiaValueState(

        KeyedStateBackend<ByteBuffer> gaiaStateBackend,

        String stateId,

        StateNamespace namespace,

        Coder<T> coder) {



      this.namespace = namespace;

      this.stateId = stateId;

      this.gaiaStateBackend = gaiaStateBackend;



      gaiaStateDescriptor = new ValueStateDescriptor<>(stateId, new CoderTypeSerializer<>(coder));

    }



    @Override

    public void write(T input) {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .update(input);

      } catch (Exception e) {

        throw new RuntimeException("Error updating state.", e);

      }

    }



    @Override

    public ValueState<T> readLater() {

      return this;

    }



    @Override

    public T read() {

      try {

        return gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .value();

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public void clear() {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .clear();

      } catch (Exception e) {

        throw new RuntimeException("Error clearing state.", e);

      }

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaValueState<?> that = (GaiaValueState<?>) o;



      return namespace.equals(that.namespace) && stateId.equals(that.stateId);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + stateId.hashCode();

      return result;

    }

  }



  private static class GaiaBagState<K, T> implements BagState<T> {



    private final StateNamespace namespace;

    private final String stateId;

    private final ListStateDescriptor<T> gaiaStateDescriptor;

    private final KeyedStateBackend<ByteBuffer> gaiaStateBackend;

    private final boolean storesVoidValues;



    GaiaBagState(

        KeyedStateBackend<ByteBuffer> gaiaStateBackend,

        String stateId,

        StateNamespace namespace,

        Coder<T> coder) {



      this.namespace = namespace;

      this.stateId = stateId;

      this.gaiaStateBackend = gaiaStateBackend;

      this.storesVoidValues = coder instanceof VoidCoder;

      this.gaiaStateDescriptor =

          new ListStateDescriptor<>(stateId, new CoderTypeSerializer<>(coder));

    }



    @Override

    public void add(T input) {

      try {

        ListState<T> partitionedState =

            gaiaStateBackend.getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor);

        if (storesVoidValues) {

          Preconditions.checkState(input == null, "Expected to a null value but was: %s", input);

          // Gaia does not allow storing null values

          // If we have null values, we use the structural null value

          input = (T) VoidCoder.of().structuralValue((Void) input);

        }

        partitionedState.add(input);

      } catch (Exception e) {

        throw new RuntimeException("Error adding to bag state.", e);

      }

    }



    @Override

    public BagState<T> readLater() {

      return this;

    }



    @Override

    @Nonnull

    public Iterable<T> read() {

      try {

        ListState<T> partitionedState =

            gaiaStateBackend.getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor);

        Iterable<T> result = partitionedState.get();

        if (storesVoidValues) {

          return () -> {

            final Iterator underlying = result.iterator();

            return new Iterator<T>() {

              @Override

              public boolean hasNext() {

                return underlying.hasNext();

              }



              @Override

              public T next() {

                // Simply move the iterator forward but ignore the value.

                // The value can be the structural null value or NULL itself,

                // if this has been restored from serialized state.

                underlying.next();

                return null;

              }

            };

          };

        }

        return result != null ? ImmutableList.copyOf(result) : Collections.emptyList();

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public ReadableState<Boolean> isEmpty() {

      return new ReadableState<Boolean>() {

        @Override

        public Boolean read() {

          try {

            Iterable<T> result =

                gaiaStateBackend

                    .getPartitionedState(

                        namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                    .get();

            return result == null;

          } catch (Exception e) {

            throw new RuntimeException("Error reading state.", e);

          }

        }



        @Override

        public ReadableState<Boolean> readLater() {

          return this;

        }

      };

    }



    @Override

    public void clear() {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .clear();

      } catch (Exception e) {

        throw new RuntimeException("Error clearing state.", e);

      }

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaBagState<?, ?> that = (GaiaBagState<?, ?>) o;



      return namespace.equals(that.namespace) && stateId.equals(that.stateId);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + stateId.hashCode();

      return result;

    }

  }



  private static class GaiaCombiningState<K, InputT, AccumT, OutputT>

      implements CombiningState<InputT, AccumT, OutputT> {



    private final StateNamespace namespace;

    private final String stateId;

    private final Combine.CombineFn<InputT, AccumT, OutputT> combineFn;

    private final ValueStateDescriptor<AccumT> gaiaStateDescriptor;

    private final KeyedStateBackend<ByteBuffer> gaiaStateBackend;



    GaiaCombiningState(

        KeyedStateBackend<ByteBuffer> gaiaStateBackend,

        String stateId,

        Combine.CombineFn<InputT, AccumT, OutputT> combineFn,

        StateNamespace namespace,

        Coder<AccumT> accumCoder) {



      this.namespace = namespace;

      this.stateId = stateId;

      this.combineFn = combineFn;

      this.gaiaStateBackend = gaiaStateBackend;



      gaiaStateDescriptor =

          new ValueStateDescriptor<>(stateId, new CoderTypeSerializer<>(accumCoder));

    }



    @Override

    public CombiningState<InputT, AccumT, OutputT> readLater() {

      return this;

    }



    @Override

    public void add(InputT value) {

      try {

        com.bff.gaia.api.common.state.ValueState<AccumT> state =

            gaiaStateBackend.getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor);



        AccumT current = state.value();

        if (current == null) {

          current = combineFn.createAccumulator();

        }

        current = combineFn.addInput(current, value);

        state.update(current);

      } catch (Exception e) {

        throw new RuntimeException("Error adding to state.", e);

      }

    }



    @Override

    public void addAccum(AccumT accum) {

      try {

        com.bff.gaia.api.common.state.ValueState<AccumT> state =

            gaiaStateBackend.getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor);



        AccumT current = state.value();

        if (current == null) {

          state.update(accum);

        } else {

          current = combineFn.mergeAccumulators(Lists.newArrayList(current, accum));

          state.update(current);

        }

      } catch (Exception e) {

        throw new RuntimeException("Error adding to state.", e);

      }

    }



    @Override

    public AccumT getAccum() {

      try {

        AccumT accum =

            gaiaStateBackend

                .getPartitionedState(

                    namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                .value();

        return accum != null ? accum : combineFn.createAccumulator();

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {

      return combineFn.mergeAccumulators(accumulators);

    }



    @Override

    public OutputT read() {

      try {

        com.bff.gaia.api.common.state.ValueState<AccumT> state =

            gaiaStateBackend.getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor);



        AccumT accum = state.value();

        if (accum != null) {

          return combineFn.extractOutput(accum);

        } else {

          return combineFn.extractOutput(combineFn.createAccumulator());

        }

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public ReadableState<Boolean> isEmpty() {

      return new ReadableState<Boolean>() {

        @Override

        public Boolean read() {

          try {

            return gaiaStateBackend

                    .getPartitionedState(

                        namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                    .value()

                == null;

          } catch (Exception e) {

            throw new RuntimeException("Error reading state.", e);

          }

        }



        @Override

        public ReadableState<Boolean> readLater() {

          return this;

        }

      };

    }



    @Override

    public void clear() {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .clear();

      } catch (Exception e) {

        throw new RuntimeException("Error clearing state.", e);

      }

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaCombiningState<?, ?, ?, ?> that = (GaiaCombiningState<?, ?, ?, ?>) o;



      return namespace.equals(that.namespace) && stateId.equals(that.stateId);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + stateId.hashCode();

      return result;

    }

  }



  private static class GaiaCombiningStateWithContext<K, InputT, AccumT, OutputT>

      implements CombiningState<InputT, AccumT, OutputT> {



    private final StateNamespace namespace;

    private final String stateId;

    private final CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn;

    private final ValueStateDescriptor<AccumT> gaiaStateDescriptor;

    private final KeyedStateBackend<ByteBuffer> gaiaStateBackend;

    private final CombineWithContext.Context context;



    GaiaCombiningStateWithContext(

        KeyedStateBackend<ByteBuffer> gaiaStateBackend,

        String stateId,

        CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn,

        StateNamespace namespace,

        Coder<AccumT> accumCoder,

        CombineWithContext.Context context) {



      this.namespace = namespace;

      this.stateId = stateId;

      this.combineFn = combineFn;

      this.gaiaStateBackend = gaiaStateBackend;

      this.context = context;



      gaiaStateDescriptor =

          new ValueStateDescriptor<>(stateId, new CoderTypeSerializer<>(accumCoder));

    }



    @Override

    public CombiningState<InputT, AccumT, OutputT> readLater() {

      return this;

    }



    @Override

    public void add(InputT value) {

      try {

        com.bff.gaia.api.common.state.ValueState<AccumT> state =

            gaiaStateBackend.getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor);



        AccumT current = state.value();

        if (current == null) {

          current = combineFn.createAccumulator(context);

        }

        current = combineFn.addInput(current, value, context);

        state.update(current);

      } catch (Exception e) {

        throw new RuntimeException("Error adding to state.", e);

      }

    }



    @Override

    public void addAccum(AccumT accum) {

      try {

        com.bff.gaia.api.common.state.ValueState<AccumT> state =

            gaiaStateBackend.getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor);



        AccumT current = state.value();

        if (current == null) {

          state.update(accum);

        } else {

          current = combineFn.mergeAccumulators(Lists.newArrayList(current, accum), context);

          state.update(current);

        }

      } catch (Exception e) {

        throw new RuntimeException("Error adding to state.", e);

      }

    }



    @Override

    public AccumT getAccum() {

      try {

        AccumT accum =

            gaiaStateBackend

                .getPartitionedState(

                    namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                .value();

        return accum != null ? accum : combineFn.createAccumulator(context);

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {

      return combineFn.mergeAccumulators(accumulators, context);

    }



    @Override

    public OutputT read() {

      try {

        com.bff.gaia.api.common.state.ValueState<AccumT> state =

            gaiaStateBackend.getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor);



        AccumT accum = state.value();

        if (accum != null) {

          return combineFn.extractOutput(accum, context);

        } else {

          return combineFn.extractOutput(combineFn.createAccumulator(context), context);

        }

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public ReadableState<Boolean> isEmpty() {

      return new ReadableState<Boolean>() {

        @Override

        public Boolean read() {

          try {

            return gaiaStateBackend

                    .getPartitionedState(

                        namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                    .value()

                == null;

          } catch (Exception e) {

            throw new RuntimeException("Error reading state.", e);

          }

        }



        @Override

        public ReadableState<Boolean> readLater() {

          return this;

        }

      };

    }



    @Override

    public void clear() {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .clear();

      } catch (Exception e) {

        throw new RuntimeException("Error clearing state.", e);

      }

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaCombiningStateWithContext<?, ?, ?, ?> that =

          (GaiaCombiningStateWithContext<?, ?, ?, ?>) o;



      return namespace.equals(that.namespace) && stateId.equals(that.stateId);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + stateId.hashCode();

      return result;

    }

  }



  private static class GaiaWatermarkHoldState<K, W extends BoundedWindow>

      implements WatermarkHoldState {



    private final TimestampCombiner timestampCombiner;

    private final Map<String, Instant> watermarkHolds;

    private final String namespaceString;

    private com.bff.gaia.api.common.state.MapState<String, Instant> watermarkHoldsState;



    public GaiaWatermarkHoldState(

        KeyedStateBackend<ByteBuffer> gaiaStateBackend,

        Map<String, Instant> watermarkHolds,

        MapStateDescriptor<String, Instant> watermarkHoldStateDescriptor,

        String stateId,

        StateNamespace namespace,

        TimestampCombiner timestampCombiner) {

      this.timestampCombiner = timestampCombiner;

      this.watermarkHolds = watermarkHolds;

      // Combines StateNamespace and stateId to generate a unique namespace for

      // watermarkHoldsState. We do not want to use Gaia's namespacing to be

      // able to recover watermark holds efficiently during recovery.

      this.namespaceString = namespace.stringKey() + stateId;

      try {

        this.watermarkHoldsState =

            gaiaStateBackend.getPartitionedState(

                VoidNamespace.INSTANCE,

                VoidNamespaceSerializer.INSTANCE,

                watermarkHoldStateDescriptor);

      } catch (Exception e) {

        throw new RuntimeException("Could not access state for watermark partition view");

      }

    }



    @Override

    public TimestampCombiner getTimestampCombiner() {

      return timestampCombiner;

    }



    @Override

    public WatermarkHoldState readLater() {

      return this;

    }



    @Override

    public ReadableState<Boolean> isEmpty() {

      return new ReadableState<Boolean>() {

        @Override

        public Boolean read() {

          try {

            return watermarkHoldsState.get(namespaceString) == null;

          } catch (Exception e) {

            throw new RuntimeException("Error reading state.", e);

          }

        }



        @Override

        public ReadableState<Boolean> readLater() {

          return this;

        }

      };

    }



    @Override

    public void add(Instant value) {

      try {

        Instant current = watermarkHoldsState.get(namespaceString);

        if (current == null) {

          watermarkHolds.put(namespaceString, value);

          watermarkHoldsState.put(namespaceString, value);

        } else {

          Instant combined = timestampCombiner.combine(current, value);

          watermarkHolds.put(namespaceString, combined);

          watermarkHoldsState.put(namespaceString, combined);

        }

      } catch (Exception e) {

        throw new RuntimeException("Error updating state.", e);

      }

    }



    @Override

    public Instant read() {

      try {

        return watermarkHoldsState.get(namespaceString);

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public void clear() {

      watermarkHolds.remove(namespaceString);

      try {

        watermarkHoldsState.remove(namespaceString);

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaWatermarkHoldState<?, ?> that = (GaiaWatermarkHoldState<?, ?>) o;



      if (!timestampCombiner.equals(that.timestampCombiner)) {

        return false;

      }

      return namespaceString.equals(that.namespaceString);

    }



    @Override

    public int hashCode() {

      int result = namespaceString.hashCode();

      result = 31 * result + timestampCombiner.hashCode();

      return result;

    }

  }



  private static class GaiaMapState<KeyT, ValueT> implements MapState<KeyT, ValueT> {



    private final StateNamespace namespace;

    private final String stateId;

    private final MapStateDescriptor<KeyT, ValueT> gaiaStateDescriptor;

    private final KeyedStateBackend<ByteBuffer> gaiaStateBackend;



    GaiaMapState(

        KeyedStateBackend<ByteBuffer> gaiaStateBackend,

        String stateId,

        StateNamespace namespace,

        Coder<KeyT> mapKeyCoder,

        Coder<ValueT> mapValueCoder) {

      this.namespace = namespace;

      this.stateId = stateId;

      this.gaiaStateBackend = gaiaStateBackend;

      this.gaiaStateDescriptor =

          new MapStateDescriptor<>(

              stateId,

              new CoderTypeSerializer<>(mapKeyCoder),

              new CoderTypeSerializer<>(mapValueCoder));

    }



    @Override

    public ReadableState<ValueT> get(final KeyT input) {

      try {

        return ReadableStates.immediate(

            gaiaStateBackend

                .getPartitionedState(

                    namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                .get(input));

      } catch (Exception e) {

        throw new RuntimeException("Error get from state.", e);

      }

    }



    @Override

    public void put(KeyT key, ValueT value) {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .put(key, value);

      } catch (Exception e) {

        throw new RuntimeException("Error put kv to state.", e);

      }

    }



    @Override

    public ReadableState<ValueT> putIfAbsent(final KeyT key, final ValueT value) {

      try {

        ValueT current =

            gaiaStateBackend

                .getPartitionedState(

                    namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                .get(key);



        if (current == null) {

          gaiaStateBackend

              .getPartitionedState(

                  namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

              .put(key, value);

        }

        return ReadableStates.immediate(current);

      } catch (Exception e) {

        throw new RuntimeException("Error put kv to state.", e);

      }

    }



    @Override

    public void remove(KeyT key) {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .remove(key);

      } catch (Exception e) {

        throw new RuntimeException("Error remove map state key.", e);

      }

    }



    @Override

    public ReadableState<Iterable<KeyT>> keys() {

      return new ReadableState<Iterable<KeyT>>() {

        @Override

        public Iterable<KeyT> read() {

          try {

            Iterable<KeyT> result =

                gaiaStateBackend

                    .getPartitionedState(

                        namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                    .keys();

            return result != null ? ImmutableList.copyOf(result) : Collections.emptyList();

          } catch (Exception e) {

            throw new RuntimeException("Error get map state keys.", e);

          }

        }



        @Override

        public ReadableState<Iterable<KeyT>> readLater() {

          return this;

        }

      };

    }



    @Override

    public ReadableState<Iterable<ValueT>> values() {

      return new ReadableState<Iterable<ValueT>>() {

        @Override

        public Iterable<ValueT> read() {

          try {

            Iterable<ValueT> result =

                gaiaStateBackend

                    .getPartitionedState(

                        namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                    .values();

            return result != null ? ImmutableList.copyOf(result) : Collections.emptyList();

          } catch (Exception e) {

            throw new RuntimeException("Error get map state values.", e);

          }

        }



        @Override

        public ReadableState<Iterable<ValueT>> readLater() {

          return this;

        }

      };

    }



    @Override

    public ReadableState<Iterable<Map.Entry<KeyT, ValueT>>> entries() {

      return new ReadableState<Iterable<Map.Entry<KeyT, ValueT>>>() {

        @Override

        public Iterable<Map.Entry<KeyT, ValueT>> read() {

          try {

            Iterable<Map.Entry<KeyT, ValueT>> result =

                gaiaStateBackend

                    .getPartitionedState(

                        namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                    .entries();

            return result != null ? ImmutableList.copyOf(result) : Collections.emptyList();

          } catch (Exception e) {

            throw new RuntimeException("Error get map state entries.", e);

          }

        }



        @Override

        public ReadableState<Iterable<Map.Entry<KeyT, ValueT>>> readLater() {

          return this;

        }

      };

    }



    @Override

    public void clear() {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .clear();

      } catch (Exception e) {

        throw new RuntimeException("Error clearing state.", e);

      }

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaMapState<?, ?> that = (GaiaMapState<?, ?>) o;



      return namespace.equals(that.namespace) && stateId.equals(that.stateId);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + stateId.hashCode();

      return result;

    }

  }



  private static class GaiaSetState<T> implements SetState<T> {



    private final StateNamespace namespace;

    private final String stateId;

    private final MapStateDescriptor<T, Boolean> gaiaStateDescriptor;

    private final KeyedStateBackend<ByteBuffer> gaiaStateBackend;



    GaiaSetState(

        KeyedStateBackend<ByteBuffer> gaiaStateBackend,

        String stateId,

        StateNamespace namespace,

        Coder<T> coder) {

      this.namespace = namespace;

      this.stateId = stateId;

      this.gaiaStateBackend = gaiaStateBackend;

      this.gaiaStateDescriptor =

          new MapStateDescriptor<>(

              stateId, new CoderTypeSerializer<>(coder), new BooleanSerializer());

    }



    @Override

    public ReadableState<Boolean> contains(final T t) {

      try {

        Boolean result =

            gaiaStateBackend

                .getPartitionedState(

                    namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                .get(t);

        return ReadableStates.immediate(result != null && result);

      } catch (Exception e) {

        throw new RuntimeException("Error contains value from state.", e);

      }

    }



    @Override

    public ReadableState<Boolean> addIfAbsent(final T t) {

      try {

        com.bff.gaia.api.common.state.MapState<T, Boolean> state =

            gaiaStateBackend.getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor);

        boolean alreadyContained = state.contains(t);

        if (!alreadyContained) {

          state.put(t, true);

        }

        return ReadableStates.immediate(!alreadyContained);

      } catch (Exception e) {

        throw new RuntimeException("Error addIfAbsent value to state.", e);

      }

    }



    @Override

    public void remove(T t) {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .remove(t);

      } catch (Exception e) {

        throw new RuntimeException("Error remove value to state.", e);

      }

    }



    @Override

    public SetState<T> readLater() {

      return this;

    }



    @Override

    public void add(T value) {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .put(value, true);

      } catch (Exception e) {

        throw new RuntimeException("Error add value to state.", e);

      }

    }



    @Override

    public ReadableState<Boolean> isEmpty() {

      return new ReadableState<Boolean>() {

        @Override

        public Boolean read() {

          try {

            Iterable<T> result =

                gaiaStateBackend

                    .getPartitionedState(

                        namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                    .keys();

            return result == null || Iterables.isEmpty(result);

          } catch (Exception e) {

            throw new RuntimeException("Error isEmpty from state.", e);

          }

        }



        @Override

        public ReadableState<Boolean> readLater() {

          return this;

        }

      };

    }



    @Override

    public Iterable<T> read() {

      try {

        Iterable<T> result =

            gaiaStateBackend

                .getPartitionedState(

                    namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

                .keys();

        return result != null ? ImmutableList.copyOf(result) : Collections.emptyList();

      } catch (Exception e) {

        throw new RuntimeException("Error read from state.", e);

      }

    }



    @Override

    public void clear() {

      try {

        gaiaStateBackend

            .getPartitionedState(

                namespace.stringKey(), StringSerializer.INSTANCE, gaiaStateDescriptor)

            .clear();

      } catch (Exception e) {

        throw new RuntimeException("Error clearing state.", e);

      }

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaSetState<?> that = (GaiaSetState<?>) o;



      return namespace.equals(that.namespace) && stateId.equals(that.stateId);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + stateId.hashCode();

      return result;

    }

  }



  /** Restores a view of the watermark holds of all keys of this partiton. */

  private void restoreWatermarkHoldsView() throws Exception {

    com.bff.gaia.api.common.state.MapState<String, Instant> mapState =

        gaiaStateBackend.getPartitionedState(

            VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, watermarkHoldStateDescriptor);

    try (Stream<ByteBuffer> keys =

        gaiaStateBackend.getKeys(watermarkHoldStateDescriptor.getName(), VoidNamespace.INSTANCE)) {

      Iterator<ByteBuffer> iterator = keys.iterator();

      while (iterator.hasNext()) {

        gaiaStateBackend.setCurrentKey(iterator.next());

        mapState.entries().forEach(entry -> watermarkHolds.put(entry.getKey(), entry.getValue()));

      }

    }

  }

}